# -*- coding: utf-8 -*-
"""
MCP服务管理器

此模块提供了MCP客户端实例的管理功能，支持单例模式和实例池管理。
"""

import asyncio
from typing import Dict, Optional

from app.config import logger
from app.services.mcp_server.mcp_client import MCPClient


class MCPManager:
    """MCP服务管理器

    提供MCP客户端实例的管理功能，确保每个MCP服务器URL对应一个客户端实例，
    支持单例模式和实例池管理。

    Attributes:
        _instances (Dict[str, MCPClient]): MCP客户端实例字典，键为服务器URL
        _lock (asyncio.Lock): 实例管理锁
    """

    def __init__(self):
        """初始化MCP服务管理器
        """
        self._instances: Dict[str, MCPClient] = {}
        self._lock = asyncio.Lock()
        logger.info("初始化MCP服务管理器")

    async def get_client(self, server_url: Optional[str] = None) -> MCPClient:
        """获取MCP客户端实例

        如果指定的服务器URL已有客户端实例，则返回该实例；否则创建新实例。

        Args:
            server_url (Optional[str]): MCP服务器URL，如果为None则使用默认URL

        Returns:
            MCPClient: MCP客户端实例
        """
        async with self._lock:
            # 使用默认URL如果未提供
            if server_url is None:
                from app.config import MCP_SERVER_URL
                server_url = MCP_SERVER_URL

            # 如果实例不存在或已断开连接，则创建新实例
            if server_url not in self._instances or not self._instances[server_url].is_connected:
                logger.info(f"为服务器URL: {server_url} 创建新的MCP客户端实例")
                client = MCPClient(server_url)
                await client.connect()
                self._instances[server_url] = client
            
            logger.debug(f"获取MCP客户端实例: {server_url}")
            return self._instances[server_url]

    async def release_client(self, server_url: Optional[str] = None) -> None:
        """释放MCP客户端实例

        断开指定服务器URL的客户端连接并从实例字典中移除。

        Args:
            server_url (Optional[str]): MCP服务器URL，如果为None则使用默认URL
        """
        async with self._lock:
            # 使用默认URL如果未提供
            if server_url is None:
                from app.config import MCP_SERVER_URL
                server_url = MCP_SERVER_URL

            if server_url in self._instances:
                logger.info(f"释放MCP客户端实例: {server_url}")
                client = self._instances.pop(server_url)
                await client.disconnect()
                logger.info(f"MCP客户端实例已释放: {server_url}")
            else:
                logger.warning(f"尝试释放不存在的MCP客户端实例: {server_url}")

    async def release_all(self) -> None:
        """释放所有MCP客户端实例

        断开所有客户端连接并清空实例字典。
        """
        async with self._lock:
            logger.info("释放所有MCP客户端实例")
            for server_url, client in list(self._instances.items()):
                try:
                    await client.disconnect()
                    logger.info(f"MCP客户端实例已释放: {server_url}")
                except Exception as e:
                    logger.error(f"释放MCP客户端实例失败: {server_url}, 错误: {str(e)}")
            self._instances.clear()
            logger.info("所有MCP客户端实例已释放")

    async def list_clients(self) -> Dict[str, bool]:
        """列出所有MCP客户端实例及其连接状态

        Returns:
            Dict[str, bool]: 客户端实例字典，键为服务器URL，值为连接状态
        """
        async with self._lock:
            result = {}
            for server_url, client in self._instances.items():
                result[server_url] = client.is_connected
            logger.debug(f"列出所有MCP客户端实例: {result}")
            return result


# 创建MCP服务管理器单例实例
mcp_manager = MCPManager()